# -*- coding: utf-8 -*-
# !/usr/bin/env python
"""
-------------------------------------------------
   File Name：     datalist
   Description :   
   Author :       lth
   date：          2023/1/3
-------------------------------------------------
   Change Activity:
                   2023/1/3 13:20: create this script
-------------------------------------------------
"""
__author__ = 'lth'

import os

from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms


class StyleGanv2Dataset(Dataset):
    def __init__(self, base_dir, mode="train"):
        super(StyleGanv2Dataset, self).__init__()

        self.mode = mode

        self.image_path_list = StyleGanv2Dataset.get_image_path_by_listdir(base_dir)

        self.T = transforms.Compose([
            transforms.RandomHorizontalFlip(p=0.5),
            transforms.ToTensor(),
            # transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
        ])

        self.AUG = transforms.Compose([
            transforms.RandomHorizontalFlip(p=0.5),
        ])

    def __len__(self):
        return len(self.image_path_list)

    def __getitem__(self, index):
        original_image = Image.open(self.image_path_list[index]).convert("RGB")
        # 32|64|128|256|512|1024
        # image_32 = original_image.resize([32,32])
        # image_64 = original_image.resize([64,64])
        # image_128 = original_image.resize([128, 128])
        image_256 = original_image.resize([256, 256])
        # image_512 = original_image.resize([512, 512])
        # image_1024 = original_image.resize([1024, 1024])

        # return self.T(image_32),self.T(image_64),self.T(image_128),self.T(image_256),self.T(image_512),self.T(image_1024)
        return self.T(image_256)

    @staticmethod
    def get_image_path_by_listdir(base_dir):
        data = []

        for image_path in os.listdir(base_dir):
            data.append(os.path.join(base_dir, image_path))

        return data
